import numpy as np
import pyopencl as cl
from time import time

rng = np.random.default_rng()
length = 5000000
a_np = rng.random(length, dtype=np.float32)
b_np = rng.random(length, dtype=np.float32)

ctx = cl.create_some_context()
queue = cl.CommandQueue(ctx)

mf = cl.mem_flags
a_g = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=a_np)
b_g = cl.Buffer(ctx, mf.READ_ONLY | mf.COPY_HOST_PTR, hostbuf=b_np)

prg = cl.Program(ctx, """
__kernel void sum(
    __global const float *a_g, __global const float *b_g, __global float *res_g)
{
  int gid = get_global_id(0);
  res_g[gid] = a_g[gid] + b_g[gid];
}
""").build()

res_g = cl.Buffer(ctx, mf.WRITE_ONLY, a_np.nbytes)
knl = prg.sum  # Use this Kernel object for repeated calls
tcl1 = time()
knl(queue, a_np.shape, None, a_g, b_g, res_g)
tcl2 = time()

# opencl while test
# while True:
#     knl(queue, a_np.shape, None, a_g, b_g, res_g)

res_np = np.empty_like(a_np)
cl.enqueue_copy(queue, res_np, res_g)

# Check on CPU with Numpy:
tnp1 = time()
np_sum = a_np + b_np
tnp2 = time()
error_np = res_np - (a_np + b_np)
print(f"Error:\n{error_np}")
print(f"Norm: {np.linalg.norm(error_np):.16e}")
assert np.allclose(res_np, a_np + b_np)

# Check on CPU with native python
res_na = np.empty_like(a_np)
tna1 = time()
for i in range(length):
    res_na[i] = a_np[i] + b_np[i]
tna2 = time()

# time
print(f"Python : {tna2 - tna1:.10f}s")
print(f"Numpy  : {tnp2 - tnp1:.10f}s")
print(f"OpenCL : {tcl2 - tcl1:.10f}s")
